from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma, FAISS
from langchain.prompts.example_selector import SemanticSimilarityExampleSelector
from langchain.prompts import FewShotPromptTemplate, PromptTemplate

# 创建嵌入模型
embeddings = HuggingFaceEmbeddings(
    model_name="sentence-transformers/nomic-embed-text",
    # model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
    model_kwargs={'device': 'cpu'},
    encode_kwargs={'normalize_embeddings': True}
)

# 示例数据
examples = [
    {"question": "什么是机器学习？", "answer": "机器学习是人工智能的一个分支，让计算机从数据中学习模式。"},
    {"question": "Python有哪些特点？", "answer": "Python简单易学、开源、有丰富的库支持。"},
    {"question": "神经网络如何工作？", "answer": "神经网络通过多层神经元处理输入数据，逐步提取特征。"},
    {"question": "如何开始学习编程？", "answer": "可以从基础语法开始，然后做小项目实践。"},
    {"question": "深度学习与机器学习的区别？", "answer": "深度学习是机器学习的子领域，使用深层神经网络。"}
]

# 创建示例选择器
example_selector = SemanticSimilarityExampleSelector.from_examples(
    examples=examples,
    embeddings=embeddings,
    vectorstore_cls=FAISS,
    k=2  # 选择2个最相似的示例
)

# 创建提示模板
example_prompt = PromptTemplate(
    input_variables=["question", "answer"],
    template="问题: {question}\n回答: {answer}\n"
)

few_shot_prompt = FewShotPromptTemplate(
    example_selector=example_selector,
    example_prompt=example_prompt,
    prefix="你是一个AI学习助手，请参考以下示例回答问题：\n\n",
    suffix="问题: {input}\n回答:",
    input_variables=["input"],
    example_separator="\n" + "="*40 + "\n"
)

# 测试示例选择
def test_similarity(query):
    selected = example_selector.select_examples({"input": query})
    print(f"查询: {query}")
    print("选择的相似示例:")
    for i, example in enumerate(selected, 1):
        print(f"{i}. 问题: {example['question']}")
        print(f"   回答: {example['answer']}")
    print("-" * 50)

# 测试不同的查询
test_queries = [
    "如何学习人工智能？",
    "Python编程难吗？",
    "什么是神经网络？"
]

for query in test_queries:
    test_similarity(query)